/*
 * DepthDataSplitter.java
 *
 * Created on Jan 28, 2013
 *
 * Copyright (c) 2012 Artur Rataj.
 *
 * This code is distributed under the terms of the GNU Library
 * General Public License, either version 3 of the license or, at
 * your option, any later version.
 */

package pl.gliwice.iitis.spatialnn.tools;

import java.awt.image.BufferedImage;
import java.awt.image.Raster;
import java.util.*;
import java.util.regex.*;
import java.io.*;
import javax.imageio.*;

/**
 * Joins a set of scenes into three left/right/depth images, by choosing horizontal
 * lines and merging them into a single image. To lower redundancy, the lines are
 * chosen so that a maximum interline space is maintained.
 * 
 * @author Artur Rataj
 */
public class DepthDataSplitter {
    /**
     * Returns all files in a given directory structure, including
     * the subdirectories.
     *
     * @param directory                 top directory in the
     *                                  hierarchy
     * @param pattern                   pattern to filter the file names,
     *                                  exluding directory names
     * @return                          list of files
     */
    public static SortedSet<String> getAllFiles(File directory, Pattern pattern) throws IOException {
// System.out.println("** at dir = " + directory.toString());
        SortedSet<String> list = new TreeSet<>();
        if(!directory.exists())
            throw new IOException("file " + directory + " not found");
        if(!directory.isDirectory())
            throw new IOException(directory + " is not a directory");
        File[] files = directory.listFiles();
        for(int i = 0; i < files.length; ++i) {
            File sub = files[i];
            if(sub.isDirectory())
                list.addAll(getAllFiles(sub, pattern));
            else {
                String name = sub.getName();
                Matcher m = pattern.matcher(name);
// System.out.println("** found = " + name);
                if(m.matches())
                    list.add(sub.getAbsolutePath());
            }
        }
        return list;
    }
    /**
     * Deletes recursively a directory.
     * 
     * From <code>http://stackoverflow.com/questions/7768071/java-delete-a-folder-content</code>.
     * 
     * @param folder 
     */
    public static void deleteDir(File folder) {
        File[] files = folder.listFiles();
        if(files!=null) { //some JVMs return null for empty dirs
            for(File f: files) {
                if(f.isDirectory()) {
                    deleteDir(f);
                } else {
                    f.delete();
                }
            }
        }
        folder.delete();
    }
    static class ImageSet {
        public List<double[]> left = new LinkedList<>();
        public List<double[]> right = new LinkedList<>();
        public List<double[]> depth = new LinkedList<>();
    }
    /**
     * Entry method of this class.
     * 
     * @param inDir directory with scene subdirectories, four images
     * per subdirectory
     * @param outDir directory where output sets are to be generated
     * @param setsNum into how many sets to scatter these
     * scenes; two sets never contain lines from the same image.
     * @param leftName name of left color images, without extension
     * @param rightName name of right color images, without extension
     * @param leftTruthName name of left truth images, without extension
     * @param rightTruthName name of right truth images, without extension
     * @param disparityMult by how much to multiply disparity
     * values from the truth images, to get disparities in the
     * color images
     */
    protected static void split(String inDir, String outDir,
            int setsNum, String leftName,
            String rightName, String leftTruthName, String rightTruthName,
            double disparityMult) throws IOException {
        File inDirFile = new File(inDir);
        SortedSet<String> images = getAllFiles(inDirFile, Pattern.compile(".*\\.png"));
        File outDirFile = new File(outDir);
        deleteDir(outDirFile);
        outDirFile.mkdir();
        ImageSet[] sets = new ImageSet[setsNum];
        int setIndex;
        for(setIndex =0; setIndex < setsNum; ++setIndex) {
            File dir = new File(outDir + "/" + setIndex);
            dir.mkdir();
            sets[setIndex] = new ImageSet();
        }
        setIndex = -1;
        String prevScene = null;
        ImageSet currSet = null;
        for(String path : images) {
            String prefix = path.substring(0, path.lastIndexOf("/"));
            String scene = prefix.substring(prefix.lastIndexOf("/") + 1);
            boolean newScene;
            if(newScene = (prevScene == null || !scene.equals(prevScene))) {
                System.out.println("Scene \"" + scene + "\"");
                prevScene = scene;
                if(++setIndex == setsNum)
                    setIndex = 0;
                currSet = sets[setIndex];
            }
            String name = path.substring(path.lastIndexOf("/") + 1,
                    path.lastIndexOf("."));
            List<double[]> lines;
            boolean depth = false;
            if(name.equals(leftName))
                lines = currSet.left;
            else if(name.equals(rightName))
                lines = currSet.right;
            else if(depth = (name.equals(leftTruthName)))
                lines = currSet.depth;
            else if(name.equals(rightTruthName))
                // read together with the left truth image
                lines = null;
            else
                throw new RuntimeException("unknown name: \"" + name + "\"");
            if(lines != null) {
                BufferedImage i = ImageIO.read(new File(prefix + "/" + name + ".png"));
                BufferedImage j;
                if(depth)
                    j = ImageIO.read(new File(prefix + "/" + rightTruthName + ".png"));
                else
                    j = null;
                Raster r = i.getRaster();
                Raster s;
                if(j != null)
                    s = j.getRaster();
                else
                    s = null;
                
                for(int y = 0; y < r.getHeight(); ++y) {
                    double[] buffer = new double[r.getWidth()];
                    for(int x = 0; x < r.getWidth(); ++x) {
                        double brightness = 0.0;
                        for(int c = 0; c < r.getNumBands(); ++c)
                            brightness += r.getSample(x, y, c);
                        if(depth) {
                            if(r.getNumBands() != 1 || s.getNumBands() != 1)
                                throw new RuntimeException("depth images expected to be grayscale");
                            brightness += s.getSample(x, y, 0);
                        }
                        brightness = brightness/(0xff*r.getNumBands());
                        if(depth)
                            brightness = disparityMult*(brightness/2.0);
                        buffer[x] = brightness;
                    }
                    lines.add(buffer);
                }
            }
            if(newScene) {
            }
        }
        for(setIndex = 0; setIndex < setsNum; ++setIndex) {
            ImageSet set = sets[setIndex];
            System.out.println("index " + setIndex + ": " +
                    "L=" + set.left.size() + " R=" + set.right.size() +
                    " D=" + set.depth.size());
        }
    }
    /**
     * Main method of this splitter.
     * 
     * @param args args[0] is the directory with scene subdirectories,
     * args[1] says, into how many sets to scatter these
     * scenes; two sets never contain lines from the same image;
     * args[2] says, 
     * args[3] contains the output directory
     */
    public static void main(String[] args) throws IOException {
        String[] args_ = {
            "images",
            "2",
            "out"
        }; args = args_;
        split(args[0], args[3], Integer.parseInt(args[1]),
                "view1", "view5", "disp1", "disp5",
                1.0/3.0);
    }
}
